def bwt(s):
    """Apply Burrows-Wheeler transform to input string."""
    assert "\0" not in s, "Input string cannot contain null character ('\0')"
    s += "\0"  # Add end of file marker
    table = sorted(s[i:] + s[:i] for i in range(len(s)))  # Table of rotations of string
    last_column = [row[-1:] for row in table]  # Last characters of each row
    return "".join(last_column)  # Convert list of characters into string


def ibwt(r, *args):
        "Inverse Burrows-Wheeler transform. args is the original index if it was not indicated by a null byte"
 
        firstCol = "".join(sorted(r))
        count = [0]*256
        byteStart = [-1]*256
        output = [""] * len(r)
        shortcut = [None]*len(r)
        #Generates shortcut lists
        for i in range(len(r)):
                shortcutIndex = ord(r[i])
                shortcut[i] = count[shortcutIndex]
                count[shortcutIndex] += 1
                shortcutIndex = ord(firstCol[i])
                if byteStart[shortcutIndex] == -1:
                        byteStart[shortcutIndex] = i
 
        localIndex = (r.index("\x00") if not args else args[0])
        for i in range(len(r)):
                #takes the next index indicated by the transformation vector
                nextByte = r[localIndex]
                output [len(r)-i-1] = nextByte
                shortcutIndex = ord(nextByte)
                #assigns localIndex to the next index in the transformation vector
                localIndex = byteStart[shortcutIndex] + shortcut[localIndex]
        return "".join(output).rstrip("\x00")

if __name__ == '__main__':
    test = 'asdfasteqta'
    out = bwt(test)
    print test
    print out
    print ibwt(out)
